"""基于Parquet格式的文件存储类，用于存储结构化数据"""
"""
Parquet文件存储实现
用于存储结构化数据，适用于模型训练数据等
"""
import os
import pandas as pd
from typing import Dict, List, Any
from fin_senti_entity_platform.data_collection.storage.base_storage import BaseStorage, StorageFactory
from fin_senti_entity_platform.utils.logger import Logger


class ParquetStorage(BaseStorage):
    """Parquet文件存储实现类"""
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化Parquet存储
        
        Args:
            config: Parquet存储配置参数
        """
        super().__init__(config)
        self.logger = Logger().get_logger(__name__)
        self.base_path = self.config.get('base_path', './data/parquet')
        
        # 确保基础路径存在
        os.makedirs(self.base_path, exist_ok=True)
        
    def _get_file_path(self, collection_name: str) -> str:
        """
        获取集合对应的文件路径
        
        Args:
            collection_name: 集合名称
        
        Returns:
            str: 文件路径
        """
        return os.path.join(self.base_path, f"{collection_name}.parquet")
    
    def save(self, data: Dict[str, Any], collection_name: str, **kwargs) -> bool:
        """
        保存单条数据到Parquet文件
        
        Args:
            data: 要保存的数据（字典格式）
            collection_name: 集合名称
            **kwargs: 其他参数
                - mode: 写入模式，'append'或'overwrite'
        """
        try:
            file_path = self._get_file_path(collection_name)
            mode = kwargs.get('mode', 'append')
            
            # 转换数据为DataFrame
            df = pd.DataFrame([data])
            
            # 检查文件是否存在
            file_exists = os.path.exists(file_path)
            
            if mode == 'append' and file_exists:
                # 追加模式
                existing_df = pd.read_parquet(file_path)
                df = pd.concat([existing_df, df], ignore_index=True)
                
            # 保存为Parquet文件
            df.to_parquet(file_path, index=False, compression='snappy')
            self.logger.info(f"成功保存数据到Parquet文件: {file_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"保存数据到Parquet文件 {collection_name} 失败: {str(e)}")
            return False
    
    def batch_save(self, data_list: List[Dict[str, Any]], collection_name: str, **kwargs) -> bool:
        """
        批量保存数据到Parquet文件
        
        Args:
            data_list: 数据列表
            collection_name: 集合名称
            **kwargs: 其他参数
                - mode: 写入模式，'append'或'overwrite'
        
        Returns:
            bool: 批量保存是否成功
        """
        try:
            if not data_list:
                return True
                
            file_path = self._get_file_path(collection_name)
            mode = kwargs.get('mode', 'append')
            
            # 转换数据为DataFrame
            df = pd.DataFrame(data_list)
            
            # 检查文件是否存在
            file_exists = os.path.exists(file_path)
            
            if mode == 'append' and file_exists:
                # 追加模式
                existing_df = pd.read_parquet(file_path)
                df = pd.concat([existing_df, df], ignore_index=True)
                
            # 保存为Parquet文件
            df.to_parquet(file_path, index=False, compression='snappy')
            self.logger.info(f"成功批量保存 {len(data_list)} 条数据到Parquet文件: {file_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"批量保存数据到Parquet文件 {collection_name} 失败: {str(e)}")
            return False
    
    def query(self, collection_name: str, query_params: Dict[str, Any], **kwargs) -> List[Dict[str, Any]]:
        """
        查询Parquet文件中的数据
        
        Args:
            collection_name: 集合名称
            query_params: 查询条件
            **kwargs: 其他参数
                - columns: 要查询的列
                - limit: 返回记录数限制
                - skip: 跳过记录数
        
        Returns:
            List[Dict[str, Any]]: 查询结果列表
        """
        try:
            file_path = self._get_file_path(collection_name)
            
            # 检查文件是否存在
            if not os.path.exists(file_path):
                self.logger.warning(f"Parquet文件 {file_path} 不存在")
                return []
                
            # 读取Parquet文件
            df = pd.read_parquet(file_path)
            
            # 应用查询条件
            for key, value in query_params.items():
                if key in df.columns:
                    df = df[df[key] == value]
                    
            # 应用列过滤
            if 'columns' in kwargs:
                columns = [col for col in kwargs['columns'] if col in df.columns]
                df = df[columns]
                
            # 应用分页
            if 'skip' in kwargs:
                df = df.iloc[kwargs['skip']:]
                
            if 'limit' in kwargs:
                df = df.iloc[:kwargs['limit']]
                
            # 转换为字典列表并返回
            results = df.to_dict('records')
            self.logger.info(f"查询Parquet文件 {collection_name}，返回 {len(results)} 条记录")
            return results
            
        except Exception as e:
            self.logger.error(f"查询Parquet文件 {collection_name} 失败: {str(e)}")
            return []
    
    def update(self, collection_name: str, query_params: Dict[str, Any], update_data: Dict[str, Any], **kwargs) -> bool:
        """
        更新Parquet文件中的数据
        
        Args:
            collection_name: 集合名称
            query_params: 查询条件
            update_data: 更新数据
            **kwargs: 其他参数
        
        Returns:
            bool: 更新是否成功
        """
        try:
            file_path = self._get_file_path(collection_name)
            
            # 检查文件是否存在
            if not os.path.exists(file_path):
                self.logger.warning(f"Parquet文件 {file_path} 不存在")
                return False
                
            # 读取Parquet文件
            df = pd.read_parquet(file_path)
            
            # 创建查询条件的掩码
            mask = pd.Series([True] * len(df))
            for key, value in query_params.items():
                if key in df.columns:
                    mask &= (df[key] == value)
                    
            # 应用更新
            for key, value in update_data.items():
                if key in df.columns:
                    df.loc[mask, key] = value
                    
            # 保存更新后的文件
            df.to_parquet(file_path, index=False, compression='snappy')
            self.logger.info(f"成功更新Parquet文件 {collection_name} 中的数据")
            return True
            
        except Exception as e:
            self.logger.error(f"更新Parquet文件 {collection_name} 失败: {str(e)}")
            return False
    
    def delete(self, collection_name: str, query_params: Dict[str, Any], **kwargs) -> bool:
        """
        删除Parquet文件中的数据
        
        Args:
            collection_name: 集合名称
            query_params: 删除条件
            **kwargs: 其他参数
        
        Returns:
            bool: 删除是否成功
        """
        try:
            file_path = self._get_file_path(collection_name)
            
            # 检查文件是否存在
            if not os.path.exists(file_path):
                self.logger.warning(f"Parquet文件 {file_path} 不存在")
                return False
                
            # 读取Parquet文件
            df = pd.read_parquet(file_path)
            
            # 创建查询条件的掩码
            mask = pd.Series([True] * len(df))
            for key, value in query_params.items():
                if key in df.columns:
                    mask &= (df[key] == value)
                    
            # 删除匹配的行
            df = df[~mask]
            
            # 保存更新后的文件
            df.to_parquet(file_path, index=False, compression='snappy')
            self.logger.info(f"成功删除Parquet文件 {collection_name} 中的数据")
            return True
            
        except Exception as e:
            self.logger.error(f"删除Parquet文件 {collection_name} 失败: {str(e)}")
            return False
    
    def close(self):
        """Parquet存储不需要保持连接，此方法为空实现"""
        pass


# 注册Parquet存储实现
StorageFactory.register_storage('parquet', ParquetStorage)